import acl
import numpy as np
import math

class AclMallocType:
    ACL_MEM_MALLOC_HUGE_FIRST = 0
    ACL_MEM_MALLOC_HUGE_ONLY = 1
    ACL_MEM_MALLOC_NORMAL_ONLY = 2
    ACL_MEM_MALLOC_HUGE_FIRST_P2P = 3
    ACL_MEM_MALLOC_HUGE_ONLY_P2P = 4
    ACL_MEM_MALLOC_NORMAL_ONLY_P2P = 5
    ACL_MEM_TYPE_LOW_BAND_WIDTH = 6
    ACL_MEM_TYPE_HIGH_BAND_WIDTH = 7


class AclMemCopyType:
    ACL_MEMCPY_HOST_TO_HOST = 0
    ACL_MEMCPY_HOST_TO_DEVICE = 1
    ACL_MEMCPY_DEVICE_TO_HOST = 2
    ACL_MEMCPY_DEVICE_TO_DEVICE = 3


class NpuNumpy:
    def __init__(self, a: np.ndarray) -> None:
        self.np_array = a
        self.np_byte_size = a.itemsize * a.size
        self.mem_size = int(math.ceil(self.np_byte_size / 32) * 32)
        self.npu_addr, ret = acl.rt.malloc(
            self.mem_size, 0
        )
        assert (ret == 0)

    def __del__(self):
        acl.rt.free(self.npu_addr)

    def npu(self):
        ret = acl.rt.memcpy(
            self.npu_addr,
            self.mem_size,
            self.np_array.ctypes.data,
            self.np_byte_size,
            AclMemCopyType.ACL_MEMCPY_HOST_TO_DEVICE,
        )
        assert (ret == 0)
        return self.npu_addr

    def cpu(self):
        ret = acl.rt.memcpy(
            self.np_array.ctypes.data,
            self.np_byte_size,
            self.npu_addr,
            self.np_byte_size,
            AclMemCopyType.ACL_MEMCPY_DEVICE_TO_HOST,
        )
        assert (ret == 0)
        return self.np_array
